#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep
from typing import Dict, List

# 导入P4Runtime库
sys.path.append(os.path.join(os.path.dirname(__file__), '../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


class ECNController:
    """ECN控制器：动态配置交换机规则和ECN阈值"""
    def __init__(self, p4info_helper, bmv2_json_path, ecn_threshold=10):
        self.p4_helper = p4info_helper
        self.bmv2_path = bmv2_json_path
        self.ecn_threshold = ecn_threshold  # 用户可配置的ECN阈值
        self.switches = {}  # 交换机连接字典

    def connect_switches(self):
        """建立与交换机的连接"""
        # 交换机配置：名称 -> (IP:端口, 设备ID)
        switch_configs = [
            ('s1', '127.0.0.1:50051', 0),
            ('s2', '127.0.0.1:50052', 1),
            ('s3', '127.0.0.1:50053', 2)
        ]

        for name, addr, dev_id in switch_configs:
            self.switches[name] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name=name,
                address=addr,
                device_id=dev_id,
                proto_dump_file=f'logs/{name}-p4runtime.log'
            )
            # 成为主控制器并加载P4程序
            self.switches[name].MasterArbitrationUpdate()
            self.switches[name].SetForwardingPipelineConfig(
                p4info=self.p4_helper.p4info,
                bmv2_json_file_path=self.bmv2_path
            )
            print(f"已连接交换机 {name}")

    def _add_ipv4_route(self, switch_name, dst_ip, prefix_len, dst_mac, out_port):
        """添加IPv4路由规则"""
        table_entry = self.p4_helper.buildTableEntry(
            table_name="IngressProcessor.ipv4_route_table",
            match_fields={"hdr.ipv4.dst": (dst_ip, prefix_len)},
            action_name="IngressProcessor.forward_ipv4",
            action_params={
                "dst_mac": dst_mac,
                "out_port": out_port
            }
        )
        self.switches[switch_name].WriteTableEntry(table_entry)
        print(f"交换机 {switch_name} 添加路由: {dst_ip}/{prefix_len} -> 端口 {out_port}")

    def _config_ecn_threshold(self, switch_name):
        """配置交换机ECN阈值（动态下发）"""
        table_entry = self.p4_helper.buildTableEntry(
            table_name="EgressProcessor.ecn_check_table",
            action_name="EgressProcessor.set_ecn",
            action_params={"queue_threshold": self.ecn_threshold},  # 传递用户配置的阈值
            default_action=True  # 设置为默认动作，对所有匹配包生效
        )
        self.switches[switch_name].WriteTableEntry(table_entry)
        print(f"交换机 {switch_name} 配置ECN阈值: {self.ecn_threshold}")

    def _setup_clone_session(self, switch_name, clone_port=255):
        """配置数据包克隆会话（将ECN标记包发送到控制器）"""
        clone_entry = self.p4_helper.buildCloneSessionEntry(
            clone_session_id=100,  # 与P4中clone_preserving_field_list的ID一致
            replicas=[{
                "egress_port": clone_port,
                "instance": 0
            }]
        )
        self.switches[switch_name].WritePREEntry(clone_entry)
        print(f"交换机 {switch_name} 配置克隆会话，目标端口 {clone_port}")

    def setup_all_rules(self):
        """配置所有交换机的路由、ECN和克隆规则"""
        # 配置s1的路由（示例拓扑：s1连接h1，与s2、s3互联）
        self._add_ipv4_route('s1', '10.0.1.1', 32, '08:00:00:00:01:11', 1)  # 到h1
        self._add_ipv4_route('s1', '10.0.2.2', 32, '08:00:00:00:02:22', 2)  # 到h2（经s2）
        self._add_ipv4_route('s1', '10.0.3.3', 32, '08:00:00:00:03:33', 3)  # 到h3（经s3）

        # 配置s2的路由
        self._add_ipv4_route('s2', '10.0.2.2', 32, '08:00:00:00:02:22', 1)  # 到h2
        self._add_ipv4_route('s2', '10.0.1.1', 32, '08:00:00:00:01:11', 2)  # 到h1（经s1）
        self._add_ipv4_route('s2', '10.0.3.3', 32, '08:00:00:00:03:33', 3)  # 到h3（经s3）

        # 配置s3的路由
        self._add_ipv4_route('s3', '10.0.3.3', 32, '08:00:00:00:03:33', 1)  # 到h3
        self._add_ipv4_route('s3', '10.0.1.1', 32, '08:00:00:00:01:11', 2)  # 到h1（经s1）
        self._add_ipv4_route('s3', '10.0.2.2', 32, '08:00:00:00:02:22', 3)  # 到h2（经s2）

        # 为所有交换机配置ECN阈值和克隆规则
        for sw_name in self.switches:
            self._config_ecn_threshold(sw_name)
            self._setup_clone_session(sw_name)

    def monitor_ecn(self):
        """监控ECN标记情况（示例：打印克隆包信息）"""
        print("\n开始监控ECN标记...（按Ctrl+C停止）")
        try:
            while True:
                # 此处可扩展为读取克隆包并解析ECN标记
                sleep(2)
        except KeyboardInterrupt:
            print("\n监控停止")


def main():
    parser = argparse.ArgumentParser(description='ECN控制器（动态配置阈值）')
    parser.add_argument('--p4info', type=str, default='build/ecn.p4.p4info.txt',
                      help='P4info文件路径')
    parser.add_argument('--bmv2-json', type=str, default='build/ecn.json',
                      help='BMv2 JSON文件路径')
    parser.add_argument('--threshold', type=int, default=10,
                      help='ECN队列长度阈值（默认10）')
    args = parser.parse_args()

    # 初始化P4信息解析器
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(args.p4info)
    # 创建控制器实例，传入用户配置的阈值
    controller = ECNController(p4info_helper, args.bmv2_json, args.threshold)

    # 启动控制器
    try:
        controller.connect_switches()
        controller.setup_all_rules()
        controller.monitor_ecn()
    except KeyboardInterrupt:
        print("\n用户中断，关闭控制器")
    finally:
        ShutdownAllSwitchConnections()


if __name__ == '__main__':
    main()